import os
import sys
import pickle
import argparse
import pandas as pd


def plot_and_save_from_history_dict(hist: dict, save_prefix: str):
    """从历史字典绘制并保存训练曲线和CSV。
    与 train_hybrid_model.py 中的实现保持一致的降级策略。
    """
    if not isinstance(hist, dict) or len(hist) == 0:
        print("⚠️ 历史为空，无法绘图。")
        return

    # 优先保存 CSV
    try:
        pd.DataFrame(hist).to_csv(f"{save_prefix}_training_history.csv", index=False)
    except Exception as e:
        print(f"⚠️ 保存训练历史CSV失败: {e}")

    # 动态导入绘图库
    try:
        import importlib
        matplotlib = importlib.import_module('matplotlib')
        matplotlib.use('Agg')
        plt = importlib.import_module('matplotlib.pyplot')
        sns = importlib.import_module('seaborn')
    except Exception as e:
        print(f"⚠️ 未安装或无法加载绘图依赖(matplotlib/seaborn)，跳过绘图。错误: {e}")
        return

    sns.set_theme(style="whitegrid")
    epochs = range(1, len(hist.get('loss', [])) + 1)

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    # Loss 曲线
    ax = axes[0]
    if 'loss' in hist:
        ax.plot(epochs, hist['loss'], label='train_loss', color='#1f77b4')
    if 'val_loss' in hist:
        ax.plot(epochs, hist['val_loss'], label='val_loss', color='#ff7f0e')
    ax.set_title('Loss over epochs')
    ax.set_xlabel('Epoch')
    ax.set_ylabel('Loss')
    ax.legend()

    # 指标曲线
    ax = axes[1]
    plotted_any_metric = False
    for metric in ['mae', 'mse', 'rmse']:
        if metric in hist:
            ax.plot(epochs, hist[metric], label=f'train_{metric}')
            plotted_any_metric = True
        if f'val_{metric}' in hist:
            ax.plot(epochs, hist[f'val_{metric}'], label=f'val_{metric}')
            plotted_any_metric = True
    if plotted_any_metric:
        ax.set_title('Metrics over epochs')
        ax.set_xlabel('Epoch')
        ax.set_ylabel('Value')
        ax.legend()
    else:
        ax.axis('off')

    plt.tight_layout()
    out_path = f"{save_prefix}_training_curves.png"
    try:
        fig.savefig(out_path, dpi=200)
        print(f"🖼️ 训练曲线已保存: {out_path}")
    except Exception as e:
        print(f"⚠️ 保存训练曲线失败: {e}")
    finally:
        plt.close(fig)


def main():
    parser = argparse.ArgumentParser(description='Plot training curves from a saved training_info.pkl')
    parser.add_argument('--info', required=True, help='Path to *_training_info.pkl')
    args = parser.parse_args()

    info_path = args.info
    if not os.path.exists(info_path):
        print(f"❌ 文件不存在: {info_path}")
        sys.exit(1)

    with open(info_path, 'rb') as f:
        training_info = pickle.load(f)

    history = training_info.get('history', {})
    save_prefix = os.path.splitext(info_path)[0].replace('_training_info', '')
    # 若上面替换不合适，退回到去除扩展名
    if not save_prefix:
        save_prefix = os.path.splitext(info_path)[0]

    plot_and_save_from_history_dict(history, save_prefix)


if __name__ == '__main__':
    main()
